# Quaternion MEKF attitude estimation simulation
# Run this cell to simulate a rigid-body rotation, generate synthetic IMU (gyro/accel/mag) data,
# run a Multiplicative EKF (MEKF) for attitude + gyro-bias estimation, and plot results.
#
# Requirements: numpy, matplotlib
# This code is self-contained and intended to be run as-is.

import numpy as np
import matplotlib.pyplot as plt

np.random.seed(1)

# ---------- helper functions (quaternion math) ----------
def skew(v):
    """Skew-symmetric matrix of a 3-vector"""
    return np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]])

def quat_mul(q1, q2):
    """Quaternion multiplication q = q1 * q2 (both as [w, x, y, z])"""
    w1, x1, y1, z1 = q1
    w2, x2, y2, z2 = q2
    return np.array([
        w1*w2 - x1*x2 - y1*y2 - z1*z2,
        w1*x2 + x1*w2 + y1*z2 - z1*y2,
        w1*y2 - x1*z2 + y1*w2 + z1*x2,
        w1*z2 + x1*y2 - y1*x2 + z1*w2
    ])

def quat_conjugate(q):
    w, x, y, z = q
    return np.array([w, -x, -y, -z])

def quat_to_rotmat(q):
    """Convert quaternion [w,x,y,z] to rotation matrix R (body <- inertial)"""
    w, x, y, z = q
    R = np.array([
        [1-2*(y*y+z*z),   2*(x*y - z*w),  2*(x*z + y*w)],
        [2*(x*y + z*w),   1-2*(x*x+z*z),  2*(y*z - x*w)],
        [2*(x*z - y*w),   2*(y*z + x*w),  1-2*(x*x+y*y)]
    ])
    return R

def normalize_quat(q):
    return q / np.linalg.norm(q)

def small_angle_quat(dtheta):
    """Convert small rotation vector (3,) to quaternion (w,x,y,z)"""
    angle = np.linalg.norm(dtheta)
    if angle < 1e-12:
        return np.array([1.0, 0.0, 0.0, 0.0])
    axis = dtheta / angle
    w = np.cos(angle/2.0)
    xyz = axis * np.sin(angle/2.0)
    return np.concatenate([[w], xyz])

def integrate_quat(q, omega, dt):
    """Integrate quaternion by angular rate omega (body rates) using first-order approx via small-angle"""
    dtheta = omega * dt
    dq = small_angle_quat(dtheta)
    q_new = quat_mul(dq, q)
    return normalize_quat(q_new)

def quat_to_euler(q):
    """Return roll, pitch, yaw (rad) from quaternion [w,x,y,z] using aerospace (ZYX yaw-pitch-roll)"""
    w, x, y, z = q
    # roll (x-axis rotation)
    sinr_cosp = 2.0*(w*x + y*z)
    cosr_cosp = 1.0 - 2.0*(x*x + y*y)
    roll = np.arctan2(sinr_cosp, cosr_cosp)
    # pitch (y-axis)
    sinp = 2.0*(w*y - z*x)
    if abs(sinp) >= 1:
        pitch = np.sign(sinp) * (np.pi/2)
    else:
        pitch = np.arcsin(sinp)
    # yaw (z-axis)
    siny_cosp = 2.0*(w*z + x*y)
    cosy_cosp = 1.0 - 2.0*(y*y + z*z)
    yaw = np.arctan2(siny_cosp, cosy_cosp)
    return np.array([roll, pitch, yaw])

# ---------- simulation parameters ----------
dt = 0.01          # sampling interval (s)
T = 20.0           # total seconds
N = int(T / dt)
t = np.arange(N) * dt

# True rotation profile (body angular velocity in body frame)
# We'll make a combined motion: slow yaw + small pitch oscillation + small roll oscillation
omega_true = np.zeros((N, 3))
omega_true[:, 2] = 0.3 * np.ones(N)                  # yaw rate ~0.3 rad/s
omega_true[:, 0] = 0.1 * np.sin(0.5 * t)             # roll oscillation
omega_true[:, 1] = 0.12 * np.cos(0.33 * t)           # pitch oscillation

# Gravity and magnetic field in inertial frame (normalized)
g_inertial = np.array([0.0, 0.0, 1.0])   # gravity unit vector (points down in body -> we assume 'down' z=+1)
m_inertial = np.array([0.2, 0.0, 0.98])  # some magnetic field vector in inertial frame
m_inertial = m_inertial / np.linalg.norm(m_inertial)

# True gyro bias (slowly varying)
gyro_bias_true = np.array([0.02, -0.015, 0.01])  # rad/s constant bias

# Sensor noise characteristics (std dev)
gyro_noise_std = 0.005    # rad/s (white noise)
gyro_bias_walk_std = 1e-4 # bias random walk std per sqrt(sec) not used for true bias here
accel_noise_std = 0.02    # unit vector noise (~g normalized)
mag_noise_std = 0.01

# Initialize true states
q_true = np.array([1.0, 0.0, 0.0, 0.0])
qs_true = np.zeros((N, 4))
qs_true[0] = q_true

for k in range(1, N):
    q_true = integrate_quat(q_true, omega_true[k-1], dt)
    qs_true[k] = q_true

# Generate sensor measurements
gyro_meas = np.zeros((N,3))
accel_meas = np.zeros((N,3))
mag_meas = np.zeros((N,3))

for k in range(N):
    qk = qs_true[k]
    Rb_i = quat_to_rotmat(qk)  # body <- inertial rotation
    # true body accel measurement is R^T * g_inertial (i.e., gravity vector expressed in body frame)
    true_accel_b = Rb_i.T @ g_inertial
    true_mag_b = Rb_i.T @ m_inertial
    # add noise and bias
    gyro_meas[k] = omega_true[k] + gyro_bias_true + np.random.randn(3) * gyro_noise_std
    accel_meas[k] = true_accel_b + np.random.randn(3) * accel_noise_std
    mag_meas[k] = true_mag_b + np.random.randn(3) * mag_noise_std

    # normalize accel and mag to emulate common practice (use direction only)
    accel_meas[k] = accel_meas[k] / np.linalg.norm(accel_meas[k])
    mag_meas[k] = mag_meas[k] / np.linalg.norm(mag_meas[k])

# ---------- MEKF estimator initialization ----------
# State: quaternion q_est (4), gyro bias b_est (3)
q_est = np.array([1.0, 0.0, 0.0, 0.0])
b_est = np.zeros(3)  # initial bias estimate

# Error-state covariance P for [delta_theta (3); delta_b (3)]
P = np.zeros((6,6))
P[0:3,0:3] = np.eye(3) * (np.deg2rad(5.0)**2)   # initial attitude uncertainty (~5 deg)
P[3:6,3:6] = np.eye(3) * (0.05**2)              # initial bias uncertainty (rad/s)

# Process noise for error-state: small-angle noise and bias random walk
# Model: delta_theta_dot = -skew(omega-b)*delta_theta - delta_b + w_theta
#        delta_b_dot = w_b
gyro_noise_var = (gyro_noise_std**2)
bias_rw_var = (gyro_bias_walk_std**2)
Q = np.zeros((6,6))
Q[0:3,0:3] = np.eye(3) * gyro_noise_var * dt   # approximate discrete process noise for angle error
Q[3:6,3:6] = np.eye(3) * bias_rw_var * dt

# Measurement noise covariances (we will use normalized vectors -> unit noise)
R_acc = np.eye(3) * (accel_noise_std**2)
R_mag = np.eye(3) * (mag_noise_std**2)

# Storage for results
q_est_store = np.zeros((N,4))
b_est_store = np.zeros((N,3))
errors = np.zeros((N,3))  # euler angle errors (rad)

# ---------- EKF loop ----------
for k in range(N):
    # 1) PREDICTION
    omega_m = gyro_meas[k]           # measured angular rate
    omega_unbiased = omega_m - b_est
    # integrate quaternion with unbiased rate
    q_est = integrate_quat(q_est, omega_unbiased, dt)

    # Discrete-time state transition Phi for error-state [dtheta; db]
    # Linearized continuous-time: d(dtheta)/dt = -skew(omega_unbiased) dtheta - I * db + w_theta
    #                         d(db)/dt = 0 + w_b
    # Discretize with first-order approx: Phi = I + F*dt
    F = np.zeros((6,6))
    F[0:3,0:3] = -skew(omega_unbiased)
    F[0:3,3:6] = -np.eye(3)
    # F[3:6,3:6] = 0 already
    Phi = np.eye(6) + F * dt
    # Propagate covariance
    P = Phi @ P @ Phi.T + Q

    # 2) UPDATE with accel (and mag) measurements
    # Use normalized accel and mag vectors as observations of R^T * g_inertial and R^T * m_inertial
    z_acc = accel_meas[k]
    z_mag = mag_meas[k]

    # predicted measurements from current q_est
    Rb_i_est = quat_to_rotmat(q_est)
    h_acc = Rb_i_est.T @ g_inertial     # expected accel direction in body
    h_mag = Rb_i_est.T @ m_inertial     # expected mag direction in body

    # build H matrix (stack accel and mag) : H = [H_acc; H_mag], each H_* is 3x6
    # H_acc = [-skew(h_acc), zeros(3)]
    H_acc = np.hstack([ -skew(h_acc), np.zeros((3,3)) ])
    H_mag = np.hstack([ -skew(h_mag), np.zeros((3,3)) ])
    H = np.vstack([H_acc, H_mag])

    # stacked measurement and predicted
    z = np.hstack([z_acc, z_mag])
    h = np.hstack([h_acc, h_mag])
    R_block = np.zeros((6,6))
    R_block[0:3,0:3] = R_acc
    R_block[3:6,3:6] = R_mag

    # Innovation
    y = z - h

    # Normalize innovation for vector measurements: optionally map to tangent by removing component
    # along predicted measurement to improve linearization; here we keep direct difference (works for small errors)

    S = H @ P @ H.T + R_block
    K = P @ H.T @ np.linalg.inv(S + 1e-12*np.eye(S.shape[0]))

    dx = K @ y  # dx = [dtheta; db]

    dtheta = dx[0:3]
    db = dx[3:6]

    # Update quaternion multiplicatively and bias additively
    dq = small_angle_quat(dtheta)
    q_est = normalize_quat(quat_mul(dq, q_est))
    b_est = b_est + db

    # Covariance update (Joseph form for numerical stability)
    I6 = np.eye(6)
    P = (I6 - K @ H) @ P @ (I6 - K @ H).T + K @ R_block @ K.T

    # store results
    q_est_store[k] = q_est
    b_est_store[k] = b_est

    # compute euler error (true - est)
    q_t = qs_true[k]
    # compute relative quaternion q_rel = q_est^{-1} * q_true  (if small, maps to error)
    q_est_conj = quat_conjugate(q_est)
    q_rel = quat_mul(q_est_conj, q_t)
    # convert to euler (approx)
    euler_true = quat_to_euler(q_t)
    euler_est = quat_to_euler(q_est)
    errors[k] = (euler_true - euler_est)

# ---------- Results and plots ----------
# convert quaternions to euler angles for plotting
euler_true = np.array([quat_to_euler(q) for q in qs_true])
euler_est = np.array([quat_to_euler(q) for q in q_est_store])

# unwrap for continuity
euler_true = np.unwrap(euler_true, axis=0)
euler_est = np.unwrap(euler_est, axis=0)

angle_names = ['Roll', 'Pitch', 'Yaw']

# 1 plot: true vs estimated Euler angles (all three on same axes)
plt.figure(figsize=(10,5))
plt.plot(t, euler_true[:,0], label='True Roll')
plt.plot(t, euler_est[:,0], label='Est Roll', linestyle='--')
plt.plot(t, euler_true[:,1], label='True Pitch')
plt.plot(t, euler_est[:,1], label='Est Pitch', linestyle='--')
plt.plot(t, euler_true[:,2], label='True Yaw')
plt.plot(t, euler_est[:,2], label='Est Yaw', linestyle='--')
plt.xlabel('Time (s)')
plt.ylabel('Angle (rad)')
plt.title('True vs Estimated Euler Angles (roll, pitch, yaw)')
plt.legend(loc='upper right', fontsize='small')
plt.grid(True)
plt.tight_layout()
plt.show()

# 2 plot: gyro bias estimates vs true bias
plt.figure(figsize=(10,4))
plt.plot(t, np.tile(gyro_bias_true[0], N), label='True bias x')
plt.plot(t, b_est_store[:,0], label='Est bias x', linestyle='--')
plt.plot(t, np.tile(gyro_bias_true[1], N), label='True bias y')
plt.plot(t, b_est_store[:,1], label='Est bias y', linestyle='--')
plt.plot(t, np.tile(gyro_bias_true[2], N), label='True bias z')
plt.plot(t, b_est_store[:,2], label='Est bias z', linestyle='--')
plt.xlabel('Time (s)')
plt.ylabel('Gyro bias (rad/s)')
plt.title('Gyro bias: true vs estimated')
plt.legend(loc='upper right', fontsize='small')
plt.grid(True)
plt.tight_layout()
plt.show()

# 3 plot: Euler angle estimation error (rad)
plt.figure(figsize=(10,4))
plt.plot(t, errors[:,0], label='Roll error')
plt.plot(t, errors[:,1], label='Pitch error')
plt.plot(t, errors[:,2], label='Yaw error')
plt.xlabel('Time (s)')
plt.ylabel('Error (rad)')
plt.title('Euler angle estimation errors (true - est)')
plt.legend(loc='upper right', fontsize='small')
plt.grid(True)
plt.tight_layout()
plt.show()

# Print RMS errors
rms_errors = np.sqrt(np.mean(errors**2, axis=0))
print("RMS angle errors (rad): Roll={:.4f}, Pitch={:.4f}, Yaw={:.4f}".format(*rms_errors))
print("RMS angle errors (deg): Roll={:.3f}°, Pitch={:.3f}°, Yaw={:.3f}°".format(*(rms_errors * 180/np.pi)))
